#include "ExecutionEngine/fpga_sim_api.h"

#include <dirent.h>
#include <sys/stat.h>

// 此文件实现 aes 的解密逻辑

// 1. Obtain plaintext, plaintext length, and key
// 2. The AES-128 packet size is 16 bytes. Check if the plaintext and key
// lengths meet the requirements. If not, complete the remaining bytes.
// 3. VCS calls the simulation environment to perform AES-128 encryption
// 4. Process output results

//===----------------------------------------------------------------------===//
// Create simulation files for simulation environment to call
//===----------------------------------------------------------------------===//
bool CreateAesSimFile(const string &sim_template_path,
                      const string &output_file_path,
                      const vector<unsigned char> &data,
                      const vector<unsigned char> &key,
                      vector<uint32_t> &result) {
  ifstream template_file(sim_template_path);
  if (!template_file) {
    cerr << "Error opening file: " << sim_template_path << endl;
    return false;
  }
  ofstream output_file(output_file_path);
  if (!output_file) {
    cerr << "Error opening file: " << output_file_path << endl;
    return false;
  }

  string line;
  bool found_end_module = false;
  vector<unsigned char> fillData(data.size() + key.size());
  auto it = fillData.begin();
  it = copy(data.begin(), data.end(), it);
  copy(key.begin(), key.end(), it);

  while (getline(template_file, line)) {
    if (!found_end_module && line.find("endmodule") != string::npos) {
      output_file
          << "always @(start_setinputdata) begin\n    @(posedge clk) begin\n";
      for (size_t i = 0; i < fillData.size(); i += 4) {
        uint32_t val = (static_cast<uint32_t>(fillData[i]) << 24) |
                       (static_cast<uint32_t>(fillData[i + 1]) << 16) |
                       (static_cast<uint32_t>(fillData[i + 2]) << 8) |
                       static_cast<uint32_t>(fillData[i + 3]);
        output_file << "        input" << dec << i / 4 << " <= 32'h" << hex
                    << setw(8) << setfill('0') << val << ";\n";
      }
      output_file << "    end\nend\n\n";
      found_end_module = true;
    }
    output_file << line << "\n";
  }
}

//===----------------------------------------------------------------------===//
// Perform simulation and obtain simulation results
//===----------------------------------------------------------------------===//
bool GetAesResultFromVCS(const string &src_files_path,
                         const string &sim_file_path,
                         vector<uint32_t> &result) {
  regex pattern(
      R"(^a = ([a-f0-9]{8}), b = ([a-f0-9]{8}), c = ([a-f0-9]{8});  d = ([a-f0-9]{8}))");
  string command = "vcs -R -full64 -debug_access+none -nolog " +
                   src_files_path + " " + sim_file_path;
  cout << "compute aes..." << endl;
  FILE *vcs_proc = popen(command.c_str(), "r");
  if (!vcs_proc) {
    cerr << "Error running VCS process" << endl;
    return false;
  }

  char buffer[1024];
  smatch match;
  string line;
  bool found_output = false;
  while (fgets(buffer, sizeof(buffer), vcs_proc) != nullptr) {
    line = buffer;
    if (regex_search(line, match, pattern)) {
      found_output = true;
      for (size_t i = 1; i <= 4; ++i) {
        // result.push_back(stoul(match[i].str(), nullptr, 16));
        uint32_t tmp = stoul(match[i].str(), nullptr, 16);
        result[i - 1] = tmp;
      }
      break;
    }
  }
  pclose(vcs_proc);
  return found_output;
}

bool aes_padded(const string &sim_template_path, const string &output_file_path,
                const string &src_files_path,
                const vector<unsigned char> &input_data,
                const vector<unsigned char> &key, vector<uint32_t> &result) {
  if (!CreateAesSimFile(sim_template_path, output_file_path, input_data, key,
                        result))
    return false;
  if (!GetAesResultFromVCS(src_files_path, output_file_path, result)) {
    cerr << "Error getting result from VCS!" << endl;
    return false;
  }
  return true;
}

//===----------------------------------------------------------------------===//
// Plaintext filling algorithm
//===----------------------------------------------------------------------===//
vector<unsigned char> aes_padding(const vector<unsigned char> &input_data) {
  // AES-128 packet size is 16 bytes
  size_t padding_len = 16 - (input_data.size() % 64);
  size_t output_length = padding_len + input_data.size();
  // FILL 0x0D
  vector<unsigned char> output_data(output_length, 0x0D);
  memcpy(output_data.data(), input_data.data(), input_data.size());
  return output_data;
}

void computeAes128(const string &sim_template_path,
                   const string &output_file_path, const string &src_files_path,
                   const vector<unsigned char> &input_data, uint8_t *key,
                   uint8_t *cipher_addr) {

  vector<vector<unsigned char>> split_input;
  vector<uint32_t> compute_result(4);

  // 1. padding plaintext
  vector<unsigned char> padded_data = aes_padding(input_data);
  vector<unsigned char> padded_key(16);
  memcpy(padded_key.data(), key, 16);
  // 2. split plaintexted by pad, 128 bits per section
  for (size_t i = 0; i < padded_data.size(); i += 16) {
    size_t end = min(i + 16, padded_data.size());
    split_input.emplace_back(padded_data.begin() + i,
                             padded_data.begin() + end);
  }
  // comput aes
  for (size_t i = 0; i < split_input.size(); ++i) {
    if (!aes_padded(sim_template_path, output_file_path, src_files_path,
                    split_input[i], padded_key, compute_result)) {
      break;
    }
  }
  if (!cipher_addr) {
    cerr << "Error cipher_addr is NULL!" << endl;
  } else {
    memcpy(cipher_addr, compute_result.data(), 20);
  }
}

//===----------------------------------------------------------------------===//
// AES encryption interface
//===----------------------------------------------------------------------===//
void fpga_aes128(uint8_t *plain_addr, uint32_t plain_len, uint8_t *key,
                 uint8_t *cipher_addr) {

  // Get environment variables
  string env_name = string(AES_SIM_PATH);
  const char *algo_path = getenv(env_name.c_str());
  if (algo_path == nullptr) {
    cerr << "Not set env variable " << env_name << "." << endl;
    exit(1);
  }
  // template.v
  string sim_template_path = string(algo_path) + string(AES_SIM_TMP_FILE);
  // *.v  verilog file
  string src_files_path = string(algo_path) + string(AES_SIM_SRC_FILE);
  // testbentch.v
  string output_file_path = string(algo_path) + string(AES_SIM_O_FILE);

  vector<unsigned char> input_data(plain_addr, plain_addr + plain_len);
  // compute aes128
  computeAes128(sim_template_path, output_file_path, src_files_path, input_data,
                key, cipher_addr);
}

//===----------------------------------------------------------------------===//
// AES decryption interface
//===----------------------------------------------------------------------===//
void fpga_aes128Crack(uint8_t *plain_addr, uint32_t plain_len, uint8_t *key,
                      uint8_t *cipher_addr) {
  // Get environment variables
  string env_name = string(AES_CRACK_SIM_PATH);
  const char *algo_path = getenv(env_name.c_str());
  if (algo_path == nullptr) {
    cerr << "Not set env variable " << env_name << "." << endl;
    exit(1);
  }
  // template.v
  string sim_template_path = string(algo_path) + string(AES_SIM_TMP_FILE);
  // *.v  verilog file
  string src_files_path = string(algo_path) + string(AES_SIM_SRC_FILE);
  // testbentch.v
  string output_file_path = string(algo_path) + string(AES_SIM_O_FILE);

  vector<unsigned char> input_data(plain_addr, plain_addr + plain_len);
  // compute aes128
  computeAes128(sim_template_path, output_file_path, src_files_path, input_data,
                key, cipher_addr);
}
